import os
import json
import tempfile
import logging
from core.annotator.export_annotations import export_annotations

# Configure logger
logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)

class AnnotationManager:
    """
    后台逻辑管理：注释内容分为工作区（transient workspace）和持久保存区（persistent storage）。
    工作区存放当前会话中新增或修改的标注，持久保存区对应 annotations.json 文件中的内容。
    提供将工作区内容合并到持久保存区等操作。
    
    每个annotation格式：
      {
         "_id": int,              # 后台自动生成的唯一ID
         "frame_file": str,
         "bbox_info": {"x": int, "y": int, "width": int, "height": int},
         "metadata_info": { ... },
         "reid_id": int           # 可重复，多条标注可能有相同值
      }
    """
    def __init__(self):
        self.workspace = {}      # video_id -> list of annotations (当前界面工作区)
        self.persistent = {}     # video_id -> list of annotations (持久保存区)
        self.video_reid_counter = {}  # video_id -> next available reid_id
        self.video_annotation_cache = {}  # 视频ID -> {帧号: 标注列表}
        self.cache_enabled = True

    # Workspace methods (transient, current session):
    def add_workspace_annotation(self, video_id: str, annotation: dict):
        """
        将新的annotation添加到工作区。如果传入的annotation中已经存在相同frame/color_id/modality组合的标注，
        则替换旧标注，否则添加新标注。不再使用reid_id字段。
        """
        if video_id not in self.workspace:
            self.workspace[video_id] = []
        
        # Check for existing annotation with the same frame/color_id/modality
        existing_idx = None
        for idx, ann in enumerate(self.workspace[video_id]):
            if (ann.get("frame") == annotation.get("frame") and
                ann.get("color_id") == annotation.get("color_id") and
                ann.get("modality") == annotation.get("modality")):
                existing_idx = idx
                break
        
        # Replace if exists, otherwise append
        if existing_idx is not None:
            # Keep the _id field if it exists in the old annotation
            if "_id" in self.workspace[video_id][existing_idx]:
                annotation["_id"] = self.workspace[video_id][existing_idx]["_id"]
            self.workspace[video_id][existing_idx] = annotation
            print(f"[INFO] 替换标注：帧:{annotation.get('frame')}, ID:{annotation.get('color_id')}, 模态:{annotation.get('modality')}")
        else:
            # Add a unique internal _id to help with identification
            annotation["_id"] = len(self.workspace.get(video_id, [])) + 1
            self.workspace[video_id].append(annotation)
            print(f"[INFO] 添加新标注：帧:{annotation.get('frame')}, ID:{annotation.get('color_id')}, 模态:{annotation.get('modality')}")
        
        # Update the cache with the new/updated annotation
        if self.cache_enabled:
            frame = annotation.get("frame")
            if frame is not None:
                if video_id not in self.video_annotation_cache:
                    self.video_annotation_cache[video_id] = {}
                    
                if frame not in self.video_annotation_cache[video_id]:
                    self.video_annotation_cache[video_id][frame] = []
                    
                # Check if this annotation exists in cache
                cache_updated = False
                for i, ann in enumerate(self.video_annotation_cache[video_id][frame]):
                    if (ann.get("color_id") == annotation.get("color_id") and
                        ann.get("modality") == annotation.get("modality")):
                        self.video_annotation_cache[video_id][frame][i] = annotation.copy()
                        cache_updated = True
                        break
                        
                if not cache_updated:
                    self.video_annotation_cache[video_id][frame].append(annotation.copy())

    def update_workspace_annotation(self, video_id: str, annotation_id: int, updated: dict):
        """
        更新工作区指定视频中 _id 对应的annotation。
        """
        if video_id not in self.workspace:
            print(f"[ERROR] 视频 {video_id} 未在工作区中找到annotations。")
            return
        for idx, ann in enumerate(self.workspace[video_id]):
            if ann.get("_id") == annotation_id:
                self.workspace[video_id][idx].update(updated)
                return
        print(f"[WARN] 未找到工作区中视频 {video_id} 的 _id {annotation_id} 的annotation。")

    def remove_workspace_annotation(self, video_id: str, annotation_id: int):
        """
        从工作区删除指定视频中 _id 对应的annotation。
        """
        if video_id not in self.workspace:
            print(f"[WARN] 视频 {video_id} 未在工作区中找到。")
            return
        before = len(self.workspace[video_id])
        self.workspace[video_id] = [ann for ann in self.workspace[video_id] if ann.get("_id") != annotation_id]
        if len(self.workspace[video_id]) == before:
            print(f"[WARN] 工作区中视频 {video_id} 未找到 _id {annotation_id} 的annotation。")
        
        # If cache is enabled, also update the cache
        if self.cache_enabled:
            for frame_number, annotations in self.video_annotation_cache.get(video_id, {}).items():
                updated_annotations = [ann for ann in annotations if ann.get("_id") != annotation_id]
                
                if len(updated_annotations) != len(annotations):
                    self.video_annotation_cache[video_id][frame_number] = updated_annotations

    def remove_annotation_by_attributes(self, video_id: str, **attributes):
        """
        根据属性（例如frame, color_id, modality）从工作区中删除匹配的标注。
        这个方法非常有用，可以确保特定属性组合的标注不会重复存在。
        
        Args:
            video_id: 视频ID
            **attributes: 用于匹配的属性，例如frame=0, color_id=1, modality="RGB"
        """
        if video_id not in self.workspace:
            return
        
        before_count = len(self.workspace[video_id])
        filtered_anns = []
        
        for ann in self.workspace[video_id]:
            # Check if all specified attributes match
            match = True
            for key, value in attributes.items():
                if ann.get(key) != value:
                    match = False
                    break
            
            # If not a match, keep the annotation
            if not match:
                filtered_anns.append(ann)
        
        # Update the workspace with annotations that don't match criteria
        self.workspace[video_id] = filtered_anns
        
        removed = before_count - len(filtered_anns)
        if removed > 0:
            attr_str = ", ".join(f"{k}={v}" for k, v in attributes.items())
            print(f"[INFO] 已从视频 {video_id} 中删除 {removed} 条符合条件 ({attr_str}) 的标注")
        
        # If cache is enabled and frame is specified, update cache
        if self.cache_enabled and "frame" in attributes:
            frame = attributes["frame"]
            if video_id in self.video_annotation_cache and frame in self.video_annotation_cache[video_id]:
                cache_annotations = self.video_annotation_cache[video_id][frame]
                filtered_cache = []
                
                for ann in cache_annotations:
                    # Check if all specified attributes match
                    match = True
                    for key, value in attributes.items():
                        if key != "frame" and ann.get(key) != value:  # Frame already matches
                            match = False
                            break
                    
                    # If not a match, keep the annotation
                    if not match:
                        filtered_cache.append(ann)
                
                # Update the cache
                self.video_annotation_cache[video_id][frame] = filtered_cache

    def get_workspace_annotations(self, video_id: str) -> list:
        """
        返回指定视频在工作区的annotation列表。
        """
        return self.workspace.get(video_id, [])

    def clear_workspace(self, video_id: str):
        """
        清空工作区中指定视频的annotations。
        """
        if video_id in self.workspace:
            self.workspace[video_id].clear()
    def annotation_exists(self, video_id: str, **attributes) -> bool:
        """
        检查是否存在具有给定属性的标注。
        """
        if video_id not in self.workspace:
            return False
            
        for ann in self.workspace[video_id]:
            match = True
            for key, value in attributes.items():
                if ann.get(key) != value:
                    match = False
                    break
            if match:
                return True
                
        return False

    # Persistent storage methods (持久保存区操作)：
    def load_persistent(self, file_path: str):
        """
        从持久保存区文件加载所有视频的annotations到persistent区，并清空工作区。
        """
        if not os.path.exists(file_path):
            print(f"[ERROR] 文件不存在: {file_path}")
            return
        with open(file_path, 'r', encoding='utf-8') as f:
            self.persistent = json.load(f)
        print(f"[INFO] Loaded persistent annotations for {len(self.persistent)} videos from: {file_path}")
        # Optionally,清空工作区
        self.workspace.clear()

    def save_persistent(self, file_path: str):
        """
        将持久保存区的annotations存入指定文件（annotations.json）。
        """
        os.makedirs(os.path.dirname(file_path), exist_ok=True)
        with open(file_path, 'w', encoding='utf-8') as f:
            json.dump(self.persistent, f, ensure_ascii=False, indent=4)
        print(f"[INFO] Persistent annotations saved to: {file_path}")

    def export_annotations(self, video_id: str, video_name: str, fps: float,
                           extract_frame_path: str, frame_count: int, modality: str,
                           output_json: str):
        """
        从持久保存区中导出指定视频的annotations到JSON文件。
        """
        ann_list = self.persistent.get(video_id, [])
        export_annotations(video_name, fps, extract_frame_path, frame_count, modality,
                           ann_list, output_json)


    # Merge workspace into persistent:
    def _strip_id(self, annotation: dict) -> dict:
        """
        Return a cleaned annotation with only essential fields.
        Ensure we don't have inconsistent ID fields in our annotations.
        """
        # Only keep essential information, excluding any ID fields except color_id
        clean_ann = {
            "frame": annotation.get("frame"),
            "bbox": annotation.get("bbox"),
            "color_id": annotation.get("color_id"),
            "modality": annotation.get("modality")
        }
        return clean_ann

    def merge_workspace(self):
        """
        将工作区中的annotations合并到持久保存区（覆盖或追加），然后清空工作区。
        持久区保存的annotation不包含工作区管理用的"_id"和"reid_id"字段。
        """
        for video_id, anns in self.workspace.items():
            if video_id not in self.persistent:
                self.persistent[video_id] = [self._strip_id(ann) for ann in anns]
            else:
                for ann in anns:
                    stripped = self._strip_id(ann)
                    # Check if a similar annotation (same frame/color_id/modality) already exists
                    exists = False
                    for i, existing in enumerate(self.persistent[video_id]):
                        if (existing.get("frame") == stripped.get("frame") and
                            existing.get("color_id") == stripped.get("color_id") and
                            existing.get("modality") == stripped.get("modality")):
                            # Replace the existing one with updated values
                            self.persistent[video_id][i] = stripped
                            exists = True
                            break
                            
                    # If no match was found, add as new
                    if not exists:
                        self.persistent[video_id].append(stripped)
        
        self.workspace.clear()
        print("[INFO] Workspace merged into persistent storage (using only essential fields).")

    def sync_annotations(self, persistent_file: str):
        """
        同步工作区的annotation到持久存储（annotations.json）。
        加载持久存储文件，将工作区内容合并（覆盖或追加不重复项），
        并保存合并后的数据回 persistent_file。
        """
        # Load existing persistent data from file, if exists.
        if os.path.exists(persistent_file):
            with open(persistent_file, 'r', encoding='utf-8') as f:
                stored = json.load(f)
            self.persistent.update(stored)
        else:
            print(f"[INFO] 持久存储文件不存在，将创建新文件：{persistent_file}")

        # Merge workspace into persistent (using _strip_id for each workspace annotation)
        for video_id, anns in self.workspace.items():
            if video_id not in self.persistent:
                self.persistent[video_id] = [self._strip_id(ann) for ann in anns]
            else:
                existing_ids = [json.dumps(ann, sort_keys=True) for ann in self.persistent[video_id]]
                for ann in anns:
                    stripped = self._strip_id(ann)
                    # Use JSON string to compare whole content.
                    if json.dumps(stripped, sort_keys=True) not in existing_ids:
                        self.persistent[video_id].append(stripped)
        # Clear the workspace after merging.
        self.workspace.clear()

        # Save updated persistent data to file.
        os.makedirs(os.path.dirname(persistent_file), exist_ok=True)
        with open(persistent_file, 'w', encoding='utf-8') as f:
            json.dump(self.persistent, f, ensure_ascii=False, indent=4)
        print(f"[INFO] 持久存储同步完成，并保存到: {persistent_file}")

    def auto_save_annotations(self):
        """Auto-save annotations for the current frame"""
        # 直接调用已经实现了模态分离的save_annotations_to_disk方法
        try:
            if not hasattr(self, 'rgb_frames') and not hasattr(self, 'ti_frames'):
                return False
                
            # 如果没有标注，跳过
            if not self.annotations:
                return True
            
            # 调用分离模态保存的方法
            result = self.save_annotations_to_disk()
            
            # 如果保存成功，返回True
            return result
        except Exception as e:
            logger.error(f"自动保存标注出错: {e}")
            import traceback
            logger.error(traceback.format_exc())
            return False

    def get_annotations_for_video(self, video_id: str, persistent_file: str) -> list:
        """
        当打开视频帧时加载对应视频的annotations：
          1. 从persistent_file中加载持久标注数据。
          2. 返回指定视频ID的标注列表（如空则返回空列表）。
        """
        self.load_persistent(persistent_file)
        return self.persistent.get(video_id, [])

    def clear_cache(self):
        """清空标注缓存"""
        self.video_annotation_cache.clear()
    
    def disable_cache(self):
        """禁用缓存"""
        self.cache_enabled = False
        
    def enable_cache(self):
        """启用缓存"""
        self.cache_enabled = True
        
    def get_cached_annotations(self, video_id: str, frame_number: int) -> list:
        """
        从缓存中获取指定视频和帧的标注，并确保返回新的副本
        """
        if not self.cache_enabled or not self.validate_cached_annotations(video_id, frame_number):
            return None
            
        # Return a deep copy to avoid unintended modifications
        cached_annotations = self.video_annotation_cache[video_id][frame_number]
        return [ann.copy() for ann in cached_annotations]
        
    def add_to_cache(self, video_id: str, frame_number: int, annotations: list):
        """
        向缓存中添加标注，确保格式一致性
        """
        if not self.cache_enabled:
            return
            
        if video_id not in self.video_annotation_cache:
            self.video_annotation_cache[video_id] = {}
            
        # Create clean copies to ensure format consistency
        cleaned_annotations = []
        for ann in annotations:
            if isinstance(ann, dict):
                # Ensure only essential fields are stored in cache
                clean_ann = self._strip_id(ann)
                cleaned_annotations.append(clean_ann)
        
        self.video_annotation_cache[video_id][frame_number] = cleaned_annotations
        print(f"[INFO] 添加到缓存: 视频{video_id}, 帧{frame_number}, {len(cleaned_annotations)}个标注")

    def validate_cached_annotations(self, video_id: str, frame_number: int) -> bool:
        """
        验证缓存中的标注是否有效
        返回：True=有效, False=无效
        """
        if not self.cache_enabled or video_id not in self.video_annotation_cache:
            return False
            
        if frame_number not in self.video_annotation_cache[video_id]:
            return False
            
        annotations = self.video_annotation_cache[video_id][frame_number]
        
        # Check we have annotations and they have required fields
        if not annotations:
            return False
            
        for ann in annotations:
            if not isinstance(ann, dict):
                return False
            if "frame" not in ann or "bbox" not in ann or "color_id" not in ann:
                return False
                
        return True

    def preload_annotations(self, video_id: str, annotation_file: str):
        """
        预加载指定视频的所有标注到缓存中
        """
        if not self.cache_enabled:
            return
            
        try:
            if os.path.exists(annotation_file):
                with open(annotation_file, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                    
                if video_id not in self.video_annotation_cache:
                    self.video_annotation_cache[video_id] = {}
                    
                if isinstance(data, dict) and "annotations" in data:
                    # 按帧号将标注分组
                    for ann in data["annotations"]:
                        if "frame" in ann:
                            frame = ann["frame"]
                            if frame not in self.video_annotation_cache[video_id]:
                                self.video_annotation_cache[video_id][frame] = []
                            self.video_annotation_cache[video_id][frame].append(ann)
                    
                    print(f"[INFO] 预加载了视频 {video_id} 的标注：{len(self.video_annotation_cache[video_id])} 帧")
                    return True
                    
        except Exception as e:
            print(f"[ERROR] 预加载标注失败 {annotation_file}: {e}")
            
        return False


